#include "obstacle_detector/pointcloud_filter.h"
#include "cuda_runtime.h"

namespace obstacle_detector
{
//计算所有解的面的三个参数
__global__ void get_all_plane_params(const int16_t *points, int iters, const int *points_indexs,
	float *plane_params, int select_num)
{
	int threadid = blockDim.x * blockIdx.x + threadIdx.x;
	if (threadid>=iters)
	{
		return;
	}

	float x_square_sum=0;
	float y_square_sum=0;
	float x_y_sum=0;
	float x_z_sum=0;
	float y_z_sum=0;
	float x_sum=0;
	float y_sum=0;
	float z_sum=0;
	float x_center;
	float y_center;
	float z_center;

	for (int i=0;i<select_num;i++)
	{
		int index=points_indexs[threadid*select_num+i];
		x_sum+=points[index*4+0]*1.0/256;
		y_sum+=points[index*4+1]*1.0/256;
		z_sum+=points[index*4+2]*1.0/256;
	}
	x_center=x_sum/select_num;
	y_center=y_sum/select_num;
	z_center=z_sum/select_num;

	for (int i=0;i<select_num;i++)
	{
		int index=points_indexs[threadid*select_num+i];
		float x_d_center=points[index*4+0]*1.0/256-x_center;
		float y_d_center=points[index*4+1]*1.0/256-y_center;
		float z_d_center=points[index*4+2]*1.0/256-z_center;
		x_square_sum+=x_d_center*x_d_center;
		y_square_sum+=y_d_center*y_d_center;
		x_y_sum+=x_d_center*y_d_center;
		x_z_sum+=x_d_center*z_d_center;
		y_z_sum+=y_d_center*z_d_center;
	}
	float d=(x_square_sum*y_square_sum - x_y_sum*x_y_sum);
	float a=(-1*y_z_sum*x_y_sum + x_z_sum*y_square_sum)/d;
	float b=(x_y_sum*-1*x_z_sum + x_square_sum*y_z_sum)/d;


	plane_params[threadid*3+0]=a;
	plane_params[threadid*3+1]=b;
	plane_params[threadid*3+2]=(z_sum-a*x_sum-b*y_sum)/select_num;
}

//计算所有面附近的点的个数
__global__ void get_best_estimate(const int16_t* points, const float* plane_params,
	int* inlier_count, int point_num,int iters, float threshold)
{
	int threadid = blockDim.x * blockIdx.x + threadIdx.x;
	if (threadid>=(point_num*iters))
	{
		return;
	}

	int point_index=threadid/iters;
	int iter_index=threadid%iters;
	float x=points[point_index*4]*1.0/256;
	float y=points[point_index*4+1]*1.0/256;
	float z=points[point_index*4+2]*1.0/256;
	float a=plane_params[iter_index*3+0];
	float b=plane_params[iter_index*3+1];
	float c=-1;
	float d=plane_params[iter_index*3+2];
	
	// printf ("diff: %f\n",x);
	if (x==0 && y==0 && z==0)
	{
		return;
	}
	// float z_predict=plane_params[iter_index*3]*x+plane_params[iter_index*3+1]*y+plane_params[iter_index*3+2];

	float dis = fabs(a*x+b*y+c*z+d)/sqrt(pow(a,2)+pow(b,2)+pow(c,2));
	// printf ("diff: %f\n",z - z_predict);
	if (dis<threshold)
	{
		atomicAdd(inlier_count+iter_index,1);
	}
}

//判断点是否在面上 并计算面上点的中心
__global__ void get_inlier_center(const int16_t *points, int point_num,
	const float* plane_params, int max_index, float *least_square_sum,
	int *inlier_index, float threshold)
{
	int threadid = blockDim.x * blockIdx.x + threadIdx.x;
	if (threadid>=(point_num))
	{
		return;
	}

	float x=points[threadid*4+0]*1.0/256;
	float y=points[threadid*4+1]*1.0/256;
	float z=points[threadid*4+2]*1.0/256;
	float a=plane_params[max_index*3+0];
	float b=plane_params[max_index*3+1];
	float c=-1;
	float d=plane_params[max_index*3+2];
	if (x==0 && y==0 && z==0)
	{
		inlier_index[threadid]=0;
		return;
	}

	float dis = fabs(a*x+b*y+c*z+d)/sqrt(pow(a,2)+pow(b,2)+pow(c,2));

	if (dis<threshold)
	{
		// atomicAdd(least_square_sum+0,x*x);
		// atomicAdd(least_square_sum+1,y*y);
		// atomicAdd(least_square_sum+2,x*y);
		// atomicAdd(least_square_sum+3,x*z);
		// atomicAdd(least_square_sum+4,y*z);
		atomicAdd(least_square_sum+0,x);
		atomicAdd(least_square_sum+1,y);
		atomicAdd(least_square_sum+2,z);
		inlier_index[threadid]=1;
	}
	else
	{
		inlier_index[threadid]=0;
	}
}
//最小二乘法需和
__global__ void get_least_square_sum(const int16_t *points, int point_num,
	const int *inlier_index, float *least_square_sum, int inlier_count)
{
	int threadid = blockDim.x * blockIdx.x + threadIdx.x;
	if (threadid>=(point_num))
	{
		return;
	}

	float x=points[threadid*4+0]*1.0/256;
	float y=points[threadid*4+1]*1.0/256;
	float z=points[threadid*4+2]*1.0/256;
	float x_center=least_square_sum[0]/inlier_count;
	float y_center=least_square_sum[1]/inlier_count;
	float z_center=least_square_sum[2]/inlier_count;

	int inlier=inlier_index[threadid];

	if (inlier==1)
	{
		float x_d_center=x-x_center;
		float y_d_center=y-y_center;
		float z_d_center=z-z_center;
		atomicAdd(least_square_sum+3,x_d_center*x_d_center);
		atomicAdd(least_square_sum+4,y_d_center*y_d_center);
		atomicAdd(least_square_sum+5,x_d_center*y_d_center);
		atomicAdd(least_square_sum+6,x_d_center*z_d_center);
		atomicAdd(least_square_sum+7,y_d_center*z_d_center);
	}
}

__global__ void get_final_points(const int16_t* points,int point_num,
	int16_t* final_points,int *final_point_num,float a,float b,float d,
	float threshold)
{
	int threadid = blockDim.x * blockIdx.x + threadIdx.x;
	if (threadid>=point_num)
	{
		return;
	}
	int16_t origin_x=points[threadid*4+0];
	int16_t origin_y=points[threadid*4+1];
	int16_t origin_z=points[threadid*4+2];
	float x=origin_x*1.0/256;
	float y=origin_y*1.0/256;
	float z=origin_z*1.0/256;
	float c=-1;
	
	// printf ("diff: %f\n",x);
	if (x==0 && y==0 && z==0)
	{
		return;
	}
	// float z_predict=plane_params[iter_index*3]*x+plane_params[iter_index*3+1]*y+plane_params[iter_index*3+2];

	float dis = fabs(a*x+b*y+c*z+d)/sqrt(pow(a,2)+pow(b,2)+pow(c,2));
	// printf ("diff: %f\n",z - z_predict);
	if (dis>threshold)
	{
		int points_cur=atomicAdd(final_point_num,1);
		final_points[points_cur*4+0]=origin_x;
		final_points[points_cur*4+1]=origin_y;
		final_points[points_cur*4+2]=origin_z;
		final_points[points_cur*4+3]=points[threadid*4+3];
	}
}



PointCloudFilter::PointCloudFilter()
{
	cudaMalloc((void **)&point_index_gpu_,sizeof(int)*RANSAC_iter_*3);
	cudaMalloc((void **)&plane_params_gpu_,sizeof(float)*RANSAC_iter_*3);
	cudaMalloc((void **)&inlier_count_gpu_,sizeof(int)*RANSAC_iter_);
	cudaMalloc((void **)&least_square_sum_gpu_,sizeof(float)*8);
	cudaMalloc((void **)&inlier_index_gpu_,sizeof(int)*200000);
	cudaMalloc((void **)&filtered_points_gpu_,sizeof(int16_t)*200000*4);
	cudaMalloc((void **)&filtered_points_num_gpu_,sizeof(int));

	inlier_count_=new int [RANSAC_iter_];
}

bool PointCloudFilter::remove_ground(int16_t* points,int &point_num)
{
	generate_point_index(point_num);

	//三个点确定一个平面
	int grid_num1=(RANSAC_iter_+1024-1)/1024;
	get_all_plane_params<<<grid_num1,1024>>>(points,RANSAC_iter_,point_index_gpu_,
		plane_params_gpu_,select_num_);

	cudaMemset(inlier_count_gpu_,0,sizeof(int)*RANSAC_iter_);
	//获取点数最多的面
	int grid_num2=(point_num*RANSAC_iter_+1024-1)/1024;
	get_best_estimate<<<grid_num2,1024>>>(points,plane_params_gpu_,
		inlier_count_gpu_,point_num,RANSAC_iter_,threshold_);
	cudaMemcpy(inlier_count_,inlier_count_gpu_,sizeof(int)*RANSAC_iter_,
		cudaMemcpyDeviceToHost);

	//获取最优解
	int max_inlier_index=std::distance(inlier_count_,std::max_element(inlier_count_,inlier_count_+RANSAC_iter_));
	// printf ("max inlier: %d,%d\n",max_inlier_index,inlier_count_[max_inlier_index]);
	if (inlier_count_[max_inlier_index]<100)//最优解是否有效
	{
		return false;
	}

	cudaMemset(least_square_sum_gpu_,0,sizeof(float)*8);
	//最优解附近的点 最小二乘估计平面
	int grid_num3=(point_num+1024-1)/1024;
	get_inlier_center<<<grid_num3,1024>>>(points,point_num,plane_params_gpu_,max_inlier_index,
		least_square_sum_gpu_,inlier_index_gpu_,threshold_);
	get_least_square_sum<<<grid_num3,1024>>>(points,point_num,inlier_index_gpu_,
		least_square_sum_gpu_,inlier_count_[max_inlier_index]);

	float least_square_sum[8];
	cudaMemcpy(least_square_sum,least_square_sum_gpu_,sizeof(float)*8,
		cudaMemcpyDeviceToHost);

	float x_sum=least_square_sum[0];
	float y_sum=least_square_sum[1];
	float z_sum=least_square_sum[2];
	float x_square_sum=least_square_sum[3];
	float y_square_sum=least_square_sum[4];
	float x_y_sum=least_square_sum[5];
	float x_z_sum=least_square_sum[6];
	float y_z_sum=least_square_sum[7];

	float d=(x_square_sum*y_square_sum - x_y_sum*x_y_sum);
	float plane[3];
	plane[0]=(-1*y_z_sum*x_y_sum + x_z_sum*y_square_sum)/d;
	plane[1]=(x_y_sum*-1*x_z_sum + x_square_sum*y_z_sum)/d;
	plane[2]=(z_sum-plane[0]*x_sum-plane[1]*y_sum)/inlier_count_[max_inlier_index];
	printf ("plane param: %f,%f,%f\n",plane[0],plane[1],plane[2]);

	cudaMemset(filtered_points_num_gpu_,0,sizeof(int));
	get_final_points<<<grid_num3,1024>>>(points,point_num,filtered_points_gpu_,
		filtered_points_num_gpu_,plane[0],plane[1],plane[2],threshold_);
	int filtered_points_num;
	cudaMemcpy(&filtered_points_num,filtered_points_num_gpu_,sizeof(int),
		cudaMemcpyDeviceToHost);
	printf ("filtered_points_num: %d,%d\n",point_num,filtered_points_num);

	cudaMemcpy(points,filtered_points_gpu_,sizeof(int16_t)*filtered_points_num*4,
		cudaMemcpyDeviceToDevice);
	point_num=filtered_points_num;
}

//生成三个一组的初始解
void PointCloudFilter::generate_point_index(int point_num)
{
	int all_index[RANSAC_iter_*3];
	for (int i=0;i<RANSAC_iter_;i++)
	{
		std::vector<int> point_index;
		while (point_index.size()<select_num_)
		{
			int selected= rand()%(point_num);
			bool duplicate=false;
			for (auto it=point_index.begin();it!=point_index.end();it++)
			{
				if (selected==(*it))
				{
					duplicate=true;
					break;
				}
			}

			if (!duplicate)
			{
				point_index.push_back(selected);
			}
		}
		for (int j=0;j<point_index.size();j++)
		{
			all_index[i*3+j]=point_index[j];
		}
	}

	cudaMemcpy(point_index_gpu_,all_index, sizeof(int)*RANSAC_iter_*3,
		cudaMemcpyHostToDevice);
	
}

}